import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

# 设置科研风格
plt.rcParams['font.family'] = 'DejaVu Sans'
plt.rcParams['font.size'] = 10
sns.set_palette("husl")

# 读取数据
df = pd.read_excel(r'D:\DeskTop\NeuTM\aa.xlsx')

# 数据结构: 第一行是指标(Cv/TD), 第一列是模型名称
# 列结构: Model | 20NG(Cv) | 20NG(TD) | NYT(Cv) | NYT(TD) | ...
models = df.iloc[1:, 0].tolist()  # 跳过第一行，获取模型名称

# 数据集和指标定义
datasets = ['20NG', 'NYT', 'NeurIPS', 'ACL', 'Wikitxt-103']
metrics = ['Cv', 'TD']

# 数据提取: 每个数据集有2列 (Cv, TD)
# 列索引: 1,2 (20NG), 3,4 (NYT), 5,6 (NeurIPS), 7,8 (ACL), 9,10 (Wikitxt-103)
data_dict = {}
for i, dataset in enumerate(datasets):
    cv_col_idx = 1 + i * 2
    td_col_idx = 2 + i * 2

    cv_values = pd.to_numeric(df.iloc[1:, cv_col_idx], errors='coerce').tolist()
    td_values = pd.to_numeric(df.iloc[1:, td_col_idx], errors='coerce').tolist()

    data_dict[dataset] = {
        'Cv': cv_values,
        'TD': td_values
    }

# 创建图形
fig, axes = plt.subplots(2, 5, figsize=(20, 8))

# 科研配色方案 - 每个模型一个颜色
colors = ['#43978F', '#9EC4BE', '#ABD0F1', '#DCE9F4', '#FBE8D5',
          '#F19685', '#F6C957', '#FFB77F', '#E56F5E']

# 创建图例的handles
from matplotlib.patches import Patch
legend_handles = [Patch(facecolor=colors[i], edgecolor='none', label=models[i])
                  for i in range(len(models))]

# 添加图例到图形顶部
fig.legend(handles=legend_handles,
           loc='upper center',
           ncol=len(models),
           frameon=True,
           fancybox=False,
           shadow=False,
           fontsize=11,
           bbox_to_anchor=(0.5, 1.02))

# 为每个数据集和指标绘制柱状图
for idx, dataset in enumerate(datasets):
    for metric_idx, metric in enumerate(metrics):
        ax = axes[metric_idx, idx]

        # 获取当前数据
        values = data_dict[dataset][metric]

        # 绘制柱状图 - 柱子紧密排列无缝隙
        x_pos = np.arange(len(models))
        bars = ax.bar(x_pos, values,
                      color=colors[:len(models)],
                      edgecolor='none',
                      width=1.0,
                      alpha=0.5)

        # 隐藏x轴标签
        ax.set_xticks(x_pos)
        ax.set_xticklabels([])

        # 设置标题和标签
        if metric_idx == 0:
            ax.set_title(f'{dataset}', fontsize=12, fontweight='bold')

        if idx == 0:
            ax.set_ylabel(f'{metric}', fontsize=11, fontweight='bold')

        # 添加网格
        ax.grid(axis='y', alpha=0.3, linestyle='--', linewidth=0.5)
        ax.set_axisbelow(True)

        # 添加数值标签
        for i, (bar, value) in enumerate(zip(bars, values)):
            if not np.isnan(value):
                ax.text(bar.get_x() + bar.get_width()/2., value,
                       f'{value:.3f}',
                       ha='center', va='bottom', fontsize=8)

        # 设置y轴范围
        valid_values = [v for v in values if not np.isnan(v)]
        if valid_values:
            ax.set_ylim(0, max(valid_values) * 1.15)

        # 美化边框
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)

# 调整布局 - 为顶部图例留出更多空间
plt.tight_layout(rect=[0, 0, 1, 0.94])

# 保存图片
plt.savefig(r'D:\DeskTop\NeuTM\results_comparison.png',
            dpi=300,
            bbox_inches='tight',
            facecolor='white',
            edgecolor='none')

print("图表已保存到: D:\\DeskTop\\NeuTM\\results_comparison.png")

# 显示图表
plt.show()
